#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
四足机器人轨迹规划程序
基于ikpy_urdf_demo.py优化，支持四条腿独立轨迹规划和参数化Tk界面
"""

import numpy as np
import ikpy.chain
import os
import sys
import time
import math
from math import pi, cos, sin
import tkinter as tk
from tkinter import ttk, messagebox, Frame
import threading
import pybullet as p
import pybullet_data

# 全局变量
running = False  # 控制仿真运行状态
leg_params = {}  # 存储四条腿的参数
leg_chains = {}  # 存储四条腿的运动链
joint_name_to_id = {}  # 关节名称到ID的映射
robot_id = None  # 机器人ID
base_link_world_pos_np = None  # 基座位置
special_markers = {}  # 末端执行器标记
trajectory_viz_points = {}  # 轨迹可视化点
trajectory_line_ids = {}  # 轨迹线ID
MAX_TRAJ_SEGMENTS = 50  # 轨迹段最大数量

# URDF文件路径
urdf_file = "robot2.urdf"

class FourLeggedRobot:
    """四足机器人类，管理机器人模型和运动学"""
    
    def __init__(self, urdf_path, use_gui=False):
        """初始化四足机器人
        
        参数:
            urdf_path: URDF文件路径
            use_gui: 是否使用GUI模式
        """
        self.urdf_path = urdf_path
        self.use_gui = use_gui
        self.robot_id = None
        self.joint_name_to_id = {}
        self.leg_chains = {}
        self.base_link_world_pos_np = None
        self.special_markers = {}
        
        # 检查URDF文件是否存在
        if not os.path.exists(urdf_path):
            raise FileNotFoundError(f"错误: URDF文件不存在: {urdf_path}")
    
    def init_pybullet(self):
        """初始化PyBullet仿真环境"""
        # 连接到PyBullet
        if self.use_gui:
            try:
                p.connect(p.GUI)
            except p.error:
                print("警告: 无法连接到GUI模式，回退到DIRECT模式")
                p.connect(p.DIRECT)
        else:
            p.connect(p.DIRECT)
            
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        p.setGravity(0, 0, -9.8)
        
        # 加载地面
        plane_id = p.loadURDF('plane.urdf')
        p.changeDynamics(plane_id, -1, lateralFriction=1)
        
        # 加载机器人
        start_pos = [0, 0, 0.5]
        start_ori = p.getQuaternionFromEuler([0, 0, 0])
        self.robot_id = p.loadURDF(self.urdf_path, start_pos, start_ori, useFixedBase=True)
        self.base_link_world_pos_np = np.array(start_pos)
        
        # 获取关节信息
        num_joints = p.getNumJoints(self.robot_id)
        print(f"机器人关节数量: {num_joints}")
        
        for i in range(num_joints):
            info = p.getJointInfo(self.robot_id, i)
            name = info[1].decode('utf-8')
            self.joint_name_to_id[name] = i
            print(f"{i}: {name}")
        
        # 设置仿真步长
        p.setTimeStep(1/240)
        
        # 设置摄像机
        if self.use_gui:
            p.resetDebugVisualizerCamera(
                cameraDistance=1.0,
                cameraYaw=50,
                cameraPitch=-30,
                cameraTargetPosition=[0, 0, 0.5]
            )
        
        return self.robot_id, self.joint_name_to_id, self.base_link_world_pos_np
    
    def create_leg_chains(self):
        """创建四条腿的运动链"""
        try:
            # 左前腿 (LF)
            print("正在从URDF文件创建左前腿运动链...")
            self.leg_chains["LF"] = ikpy.chain.Chain.from_urdf_file(
                self.urdf_path,
                base_elements=["base_link", "A1", "A1_Link", "A2", "A2_Link", "A3", "A3_Link", "A4"],
                active_links_mask=[False, True, True, True, False],
                name="left_front_leg"
            )
            
            # 右前腿 (RF)
            print("正在从URDF文件创建右前腿运动链...")
            self.leg_chains["RF"] = ikpy.chain.Chain.from_urdf_file(
                self.urdf_path,
                base_elements=["base_link", "B1", "B1_Link", "B2", "B2_Link", "B3", "B3_Link", "B4"],
                active_links_mask=[False, True, True, True, False],
                name="right_front_leg"
            )
            
            # 左后腿 (LB)
            print("正在从URDF文件创建左后腿运动链...")
            self.leg_chains["LB"] = ikpy.chain.Chain.from_urdf_file(
                self.urdf_path,
                base_elements=["base_link", "C1", "C1_Link", "C2", "C2_Link", "C3", "C3_Link", "C4"],
                active_links_mask=[False, True, True, True, False],
                name="left_back_leg"
            )
            
            # 右后腿 (RB)
            print("正在从URDF文件创建右后腿运动链...")
            self.leg_chains["RB"] = ikpy.chain.Chain.from_urdf_file(
                self.urdf_path,
                base_elements=["base_link", "D1", "D1_Link", "D2", "D2_Link", "D3", "D3_Link", "D4"],
                active_links_mask=[False, True, True, True, False],
                name="right_back_leg"
            )
            
            # 打印链信息
            for leg_id, chain in self.leg_chains.items():
                print(f"链名称: {chain.name} ({leg_id})")
                print(f"链长度: {len(chain.links)}")
                for i, link in enumerate(chain.links):
                    print(f"  链接 {i}: {link.name}")
            
            return self.leg_chains
            
        except Exception as e:
            print(f"创建运动链时出错: {e}")
            raise
    
    def create_end_effector_markers(self):
        """为每条腿创建末端执行器标记"""
        for leg_id in self.leg_chains.keys():
            # 创建红色球体可视化目标点
            special_marker_visual = p.createVisualShape(
                p.GEOM_SPHERE, 
                radius=0.01, 
                rgbaColor=[1, 0, 0, 1]
            )
            
            # 初始位置设置在远处
            special_marker_id = p.createMultiBody(
                baseVisualShapeIndex=special_marker_visual,
                basePosition=[100, 100, 100]
            )
            
            self.special_markers[leg_id] = special_marker_id
            trajectory_viz_points[leg_id] = []
            trajectory_line_ids[leg_id] = []
        
        return self.special_markers
    
    def reset_leg_positions(self):
        """重置所有腿到初始位置"""
        for leg_id, chain in self.leg_chains.items():
            zero_joint_angles = np.zeros(len(chain.links))
            self.apply_joint_angles(chain, zero_joint_angles)
    
    def apply_joint_angles(self, chain, angles):
        """应用关节角度到机器人
        参数:
            chain: 运动链
            angles: 关节角度数组
        """
        for i, mask in enumerate(chain.active_links_mask):
            if mask:
                name = chain.links[i].name
                if name in self.joint_name_to_id:
                    p.setJointMotorControl2(
                        self.robot_id,
                        self.joint_name_to_id[name],
                        p.POSITION_CONTROL,
                        targetPosition=angles[i],
                        force=500
                    )
    
    def solve_ik_and_apply(self, leg_id, target_pos_base_link, init_pos=None):
        """IK求解并应用关节角度
        
        参数:
            leg_id: 腿的ID (LF, RF, LB, RB)
            target_pos_base_link: 目标位置 (基座坐标系)
            init_pos: 初始猜测位置
            
        返回:
            solution: 关节角度解
            error: 位置误差
        """
        chain = self.leg_chains[leg_id]
        
        # 如果没有提供初始猜测，使用零向量
        if init_pos is None:
            init_pos = np.zeros(len(chain.links))
        
        # 求解IK
        solution = chain.inverse_kinematics(
            target_pos_base_link,
            initial_position=init_pos,
            max_iter=1000
        )
        
        # 应用解到机器人
        self.apply_joint_angles(chain, solution)
        
        # 计算误差
        fk = chain.forward_kinematics(solution)
        pos_base_link = fk[:3, 3]
        error = np.linalg.norm(pos_base_link - target_pos_base_link)
        
        return solution, error
    
    def update_marker_position(self, leg_id, target_pos_base_link):
        """更新标记位置
        
        参数:
            leg_id: 腿的ID
            target_pos_base_link: 目标位置 (基座坐标系)
        """
        # 转换到世界坐标系
        world_marker_pos = (target_pos_base_link + self.base_link_world_pos_np).tolist()
        
        # 更新标记位置
        p.resetBasePositionAndOrientation(
            self.special_markers[leg_id], 
            world_marker_pos, 
            [0, 0, 0, 1]
        )
        
        # 更新轨迹可视化
        if len(trajectory_viz_points[leg_id]) > 0:
            prev_ee_world_pos = trajectory_viz_points[leg_id][-1]
            line_id = p.addUserDebugLine(
                prev_ee_world_pos, 
                world_marker_pos,
                lineColorRGB=[0.2, 0.2, 1.0],
                lineWidth=1.5
            )
            trajectory_line_ids[leg_id].append(line_id)
        
        trajectory_viz_points[leg_id].append(world_marker_pos)
        
        # 限制轨迹段数量
        if len(trajectory_line_ids[leg_id]) > MAX_TRAJ_SEGMENTS:
            p.removeUserDebugItem(trajectory_line_ids[leg_id].pop(0))
            trajectory_viz_points[leg_id].pop(0)


class TrajectoryPlanner:
    """轨迹规划器类，生成轨迹点"""
    
    def __init__(self):
        """初始化轨迹规划器"""
        # 默认参数
        self.default_params = {
            "LF": {
                "center_x": -0.047, "center_y": 0.163, "base_z": -0.242,
                "long_axis": 0.075, "short_axis": 0.0, "height": 0.05,
                "direction": "Y", "phase_offset": 0.0
            },
            "RF": {
                "center_x": 0.047, "center_y": 0.163, "base_z": -0.242,
                "long_axis": 0.075, "short_axis": 0.0, "height": 0.05,
                "direction": "Y", "phase_offset": 0.5
            },
            "LB": {
                "center_x": -0.047, "center_y": -0.163, "base_z": -0.242,
                "long_axis": 0.075, "short_axis": 0.0, "height": 0.05,
                "direction": "Y", "phase_offset": 0.5
            },
            "RB": {
                "center_x": 0.047, "center_y": -0.163, "base_z": -0.242,
                "long_axis": 0.075, "short_axis": 0.0, "height": 0.05,
                "direction": "Y", "phase_offset": 0.0
            }
        }
        
        # 当前参数 (深拷贝默认参数)
        self.current_params = {}
        for leg_id, params in self.default_params.items():
            self.current_params[leg_id] = params.copy()
    
    def generate_trajectory(self, leg_id, time_phase, speed=1.0):
        """生成轨迹点
        
        参数:
            leg_id: 腿的ID
            time_phase: 时间相位 (0-1)
            speed: 速度因子
            
        返回:
            target_pos: 目标位置 (x, y, z)
        """
        params = self.current_params[leg_id]
        status = 0 
        # 计算腿部相位
        leg_phase = (time_phase * speed + params["phase_offset"]) % 1.0
        
        # 计算椭圆轨迹点
        if 0 <= leg_phase < 0.5:  # 摆动相
            swing_progress = leg_phase / 0.5
            angle = pi - swing_progress * pi
            status = 0

            if params["direction"] == "X":
                x_offset = params["long_axis"] * cos(angle)
                y_offset = params["short_axis"] * sin(angle) if params["short_axis"] > 0 else 0
            else:  # "Y"方向
                x_offset = params["short_axis"] * sin(angle) if params["short_axis"] > 0 else 0
                y_offset = params["long_axis"] * cos(angle)
                
            z_offset = params["height"] * sin(angle)
        else:  # 支撑相
            return_progress = (leg_phase - 0.5) / 0.5
            status = 1
            
            if params["direction"] == "X":
                x_offset = -params["long_axis"] + 2 * params["long_axis"] * return_progress
                y_offset = 0
            else:  # "Y"方向
                x_offset = 0
                y_offset = -params["long_axis"] + 2 * params["long_axis"] * return_progress
                
            z_offset = 0
        
        # 计算最终位置
        x = params["center_x"] + x_offset
        y = params["center_y"] + y_offset
        z = params["base_z"] + z_offset
        
        return np.array([x, y, z]),status


class TkInterface:
    """Tk界面类，提供用户界面"""
    
    def __init__(self, robot, planner):
        """初始化Tk界面
        
        参数:
            robot: FourLeggedRobot实例
            planner: TrajectoryPlanner实例
        """
        self.robot = robot
        self.planner = planner
        self.current_angles = {}
        
        # 为每条腿初始化当前角度
        for leg_id in robot.leg_chains.keys():
            self.current_angles[leg_id] = np.zeros(len(robot.leg_chains[leg_id].links))
        
        # 创建主窗口
        self.root = tk.Tk()
        self.root.title("四足机器人轨迹规划控制面板")
        self.root.geometry("800x600")
        
        # 创建控制变量
        self.create_control_variables()
        
        # 创建界面组件
        self.create_widgets()
        
        # 绑定关闭事件
        self.root.protocol("WM_DELETE_WINDOW", self.on_closing)
    
    def create_control_variables(self):
        """创建控制变量"""
        # 全局控制变量
        self.running_var = tk.BooleanVar(value=False)
        self.speed_var = tk.DoubleVar(value=0.5)
        
        # 腿部参数变量
        self.leg_vars = {}
        
        for leg_id, params in self.planner.current_params.items():
            self.leg_vars[leg_id] = {
                "center_x": tk.DoubleVar(value=params["center_x"]),
                "center_y": tk.DoubleVar(value=params["center_y"]),
                "base_z": tk.DoubleVar(value=params["base_z"]),
                "long_axis": tk.DoubleVar(value=params["long_axis"]),
                "short_axis": tk.DoubleVar(value=params["short_axis"]),
                "height": tk.DoubleVar(value=params["height"]),
                "direction": tk.StringVar(value=params["direction"])
            }
    
    def create_widgets(self):
        """创建界面组件"""
        # 创建主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 创建全局控制区域
        global_frame = ttk.LabelFrame(main_frame, text="全局控制", padding="5")
        global_frame.pack(fill=tk.X, pady=5)
        
        # 启动/停止按钮
        ttk.Checkbutton(
            global_frame, 
            text="启动/停止", 
            variable=self.running_var,
            command=self.toggle_running
        ).grid(row=0, column=0, padx=5, pady=5)
        
        # 速度控制
        ttk.Label(global_frame, text="速度:").grid(row=0, column=1, padx=5, pady=5)
        speed_scale = ttk.Scale(
            global_frame,
            from_=0.1,
            to=2.0,
            orient=tk.HORIZONTAL,
            variable=self.speed_var,
            length=200
        )
        speed_scale.grid(row=0, column=2, padx=5, pady=5)
        
        # 重置按钮
        ttk.Button(
            global_frame,
            text="重置参数",
            command=self.reset_params
        ).grid(row=0, column=3, padx=5, pady=5)
        
        # 创建腿部控制区域
        legs_frame = ttk.Frame(main_frame)
        legs_frame.pack(fill=tk.BOTH, expand=True, pady=5)
        
        # 为每条腿创建控制面板
        leg_names = {"LF": "左前腿", "RF": "右前腿", "LB": "左后腿", "RB": "右后腿"}
        
        for i, (leg_id, leg_name) in enumerate(leg_names.items()):
            row, col = divmod(i, 2)
            self.create_leg_control(legs_frame, leg_id, leg_name, row, col)
    
    def create_leg_control(self, parent, leg_id, leg_name, row, col):
        """创建腿部控制面板
        
        参数:
            parent: 父容器
            leg_id: 腿的ID
            leg_name: 腿的名称
            row: 行索引
            col: 列索引
        """
        # 创建腿部框架
        leg_frame = ttk.LabelFrame(parent, text=f"{leg_name} ({leg_id})", padding="5")
        leg_frame.grid(row=row, column=col, padx=5, pady=5, sticky="nsew")
        
        # 设置网格权重
        parent.grid_columnconfigure(0, weight=1)
        parent.grid_columnconfigure(1, weight=1)
        parent.grid_rowconfigure(0, weight=1)
        parent.grid_rowconfigure(1, weight=1)
        
        # 参数标签和输入框
        params = [
            ("中心X:", "center_x", -0.5, 0.5),
            ("中心Y:", "center_y", -0.5, 0.5),
            ("基础Z:", "base_z", -0.5, 0.0),
            ("长半轴:", "long_axis", 0.0, 0.2),
            ("短半轴:", "short_axis", 0.0, 0.2),
            ("高度:", "height", 0.0, 0.1)
        ]
        
        for i, (label_text, param_name, min_val, max_val) in enumerate(params):
            ttk.Label(leg_frame, text=label_text).grid(row=i, column=0, sticky="w", padx=5, pady=2)
            
            # 创建输入框和滑块组合
            param_frame = ttk.Frame(leg_frame)
            param_frame.grid(row=i, column=1, sticky="ew", padx=5, pady=2)
            
            # 输入框
            entry = ttk.Entry(
                param_frame,
                textvariable=self.leg_vars[leg_id][param_name],
                width=8
            )
            entry.pack(side=tk.LEFT)
            
            # 滑块
            scale = ttk.Scale(
                param_frame,
                from_=min_val,
                to=max_val,
                orient=tk.HORIZONTAL,
                variable=self.leg_vars[leg_id][param_name],
                length=150
            )
            scale.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5)
        
        # 方向选择
        ttk.Label(leg_frame, text="方向:").grid(row=len(params), column=0, sticky="w", padx=5, pady=2)
        direction_frame = ttk.Frame(leg_frame)
        direction_frame.grid(row=len(params), column=1, sticky="w", padx=5, pady=2)
        
        ttk.Radiobutton(
            direction_frame,
            text="X",
            variable=self.leg_vars[leg_id]["direction"],
            value="X"
        ).pack(side=tk.LEFT, padx=5)
        
        ttk.Radiobutton(
            direction_frame,
            text="Y",
            variable=self.leg_vars[leg_id]["direction"],
            value="Y"
        ).pack(side=tk.LEFT, padx=5)
    
    def update_planner_params(self):
        """更新规划器参数"""
        for leg_id, vars_dict in self.leg_vars.items():
            for param_name, var in vars_dict.items():
                self.planner.current_params[leg_id][param_name] = var.get()
    
    def toggle_running(self):
        """切换运行状态"""
        global running
        running = self.running_var.get()
    
    def reset_params(self):
        """重置参数到默认值"""
        # 重置规划器参数
        for leg_id, params in self.planner.default_params.items():
            self.planner.current_params[leg_id] = params.copy()
        
        # 更新界面变量
        for leg_id, params in self.planner.current_params.items():
            for param_name, value in params.items():
                if param_name in self.leg_vars[leg_id]:
                    self.leg_vars[leg_id][param_name].set(value)
    
    def update_ui(self):
        """更新界面"""
        # 在这里可以添加界面更新逻辑
        pass
    
    def on_closing(self):
        """关闭窗口处理"""
        global running
        running = False
        self.root.destroy()
    
    def start(self):
        """启动界面"""
        # 启动主循环
        self.root.mainloop()


def simulation_thread(robot, planner, ui):
    """仿真线程
    
    参数:
        robot: FourLeggedRobot实例
        planner: TrajectoryPlanner实例
        ui: TkInterface实例
    """
    global running
    
    # 当前角度字典
    current_angles = {}
    for leg_id in robot.leg_chains.keys():
        current_angles[leg_id] = np.zeros(len(robot.leg_chains[leg_id].links))
    
    # 摄像机参数
    camera_distance = 1.0
    camera_yaw = 50
    camera_pitch = -30
    camera_target = [0, 0, 0.5]
    
    # 主循环
    while True:
        try:
            # 检查界面是否关闭
            if not hasattr(ui, 'root') or not ui.root.winfo_exists():
                break
            
            # 更新规划器参数
            ui.update_planner_params()
            
            # 处理键盘事件 - 摄像头控制
            keys = p.getKeyboardEvents()
            if p.B3G_LEFT_ARROW in keys and (keys[p.B3G_LEFT_ARROW] & p.KEY_IS_DOWN):
                camera_yaw -= 2  # 左键减小yaw值，摄像头向左转
            if p.B3G_RIGHT_ARROW in keys and (keys[p.B3G_RIGHT_ARROW] & p.KEY_IS_DOWN):
                camera_yaw += 2  # 右键增加yaw值，摄像头向右转
            
            # 更新摄像机位置
            if robot.use_gui:
                p.resetDebugVisualizerCamera(
                    cameraDistance=camera_distance,
                    cameraYaw=camera_yaw,
                    cameraPitch=camera_pitch,
                    cameraTargetPosition=camera_target
                )
            
            # 如果运行标志为真，执行轨迹规划
            if running:
                # 获取当前时间相位
                current_time_phase = (time.time()) % 1.0
                speed = ui.speed_var.get()
                
                # 为每条腿生成轨迹
                for leg_id in robot.leg_chains.keys():
                    # 生成目标位置
                    target_pos = planner.generate_trajectory(leg_id, current_time_phase, speed)
                    
                    # 求解IK并应用
                    solution, error = robot.solve_ik_and_apply(
                        leg_id, 
                        target_pos, 
                        init_pos=current_angles[leg_id]
                    )
                    
                    # 更新当前角度
                    current_angles[leg_id] = solution
                    
                    # 更新标记位置
                    robot.update_marker_position(leg_id, target_pos)
            
            # 步进仿真
            p.stepSimulation()
            time.sleep(1/240)
            
        except KeyboardInterrupt:
            print("仿真线程被中断")
            break
        except Exception as e:
            print(f"仿真线程出错: {e}")
            break
    
    # 断开PyBullet连接
    p.disconnect()


def main():
    """主函数"""
    import argparse
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="四足机器人轨迹规划程序")
    args = parser.parse_args()
    
    try:
        # GUI模式
        # 创建四足机器人
        robot = FourLeggedRobot(urdf_file, use_gui=True)
        
        # 初始化PyBullet
        global robot_id, joint_name_to_id, base_link_world_pos_np
        robot_id, joint_name_to_id, base_link_world_pos_np = robot.init_pybullet()
        
        # 创建运动链
        global leg_chains
        leg_chains = robot.create_leg_chains()
        
        # 创建末端执行器标记
        global special_markers
        special_markers = robot.create_end_effector_markers()
        
        # 重置腿部位置
        robot.reset_leg_positions()
        
        # 创建轨迹规划器
        planner = TrajectoryPlanner()
        
        # 创建Tk界面
        ui = TkInterface(robot, planner)
        
        # 启动仿真线程
        sim_thread = threading.Thread(
            target=simulation_thread,
            args=(robot, planner, ui),
            daemon=True
        )
        sim_thread.start()
        
        # 启动界面
        ui.start()
        
    except KeyboardInterrupt:
        print("程序被中断")
    except Exception as e:
        print(f"程序出错: {e}")
    finally:
        # 确保断开PyBullet连接
        if p.isConnected():
            p.disconnect()


if __name__ == "__main__":
    main()
